# -*- coding: utf-8 -*-
# created on 2016/3/10

from sympy import Add, solveset, Eq, lcm, latex, Mul
from numpy import argsort, array
from collections import defaultdict
from mathsolver.functions.sympy_utils import get_pow, get_all_child
from sympy import sympify

ZHENGLI = "移项："
HUAJIAN = "化简："
MULTIPLE = "使%s的系数相等："
GET = "由式(%d)得："
PLUG = "将等式(%d)代入等式(%d)，得："
PLUG_2 = "将解代入等式(%d)，得："
ADD = "(%d) + (%d)得："
MINUS = "(%d) - (%d)得："
JIEDE = ""
zero = sympify("0")


class Step2Equations(object):
    eq_num = 2

    def __init__(self, eqs, syms, index=1):
        self.init_index = index
        self.eqs = eqs
        self.symbols = syms
        self.steps = []
        self.index = index
        self.answer = defaultdict(list)

    def check(self):
        free_symbols = set()
        for eq in self.eqs:
            free_symbols.update(eq.free_symbols)
        if len(free_symbols) == self.eq_num:
            self.symbols = list(free_symbols)

    @staticmethod
    def factor(f):
        factor_expr = f.factor()
        if factor_expr != f and factor_expr.is_Mul:
            args = [arg for arg in factor_expr.args if arg.is_Number]
            if len(args) == 1:
                return f / args[0]
        return f

    def get_coeff(self, f):  # 这里的get_coeff还需要进一步修改
        coeff = {}
        for symb in self.symbols:
            _, h = f.as_independent(symb)
            if h.is_Symbol:
                coeff[h] = 1
            elif h.is_Mul and (h.args[1].is_Symbol or h.args[1].is_Pow):
                # coeff[h.args[1]] = h.args[0]
                coeff[symb] = h / symb  # xj20160906
            elif h.is_Pow and h.args[0].is_Symbol:
                coeff[h] = 1

        return coeff

    def get_answer(self, eq, sym, known):
        if len(eq.free_symbols) > 1:
            ans = solveset(eq, sym)
        else:
            ans = solveset(eq, sym)

        if ans.is_FiniteSet and ans.is_EmptySet:
            raise ValueError("No solution")
        elif ans.is_FiniteSet:
            for arg in ans.args:
                if known:
                    for s, a in known.items():
                        self.answer[s].append(a)
                self.answer[sym].append(arg)

            if sym in self.symbols:
                self.symbols.remove(sym)
        else:
            raise ValueError("Unexpected Situation %s!" % ans)

        return ans

    def check_simple(self, eq):
        if len(eq.free_symbols) == 1:
            sym = eq.free_symbols.pop()
            if len(self.answer) > 0:
                known_sym, known_solve = zip(*self.answer.items())
                self.answer = defaultdict(list)
                for known in zip(*known_solve):
                    known = dict([(s, a) for s, a in zip(known_sym, known)])
                    self.get_answer(eq, sym, known)
            else:
                self.get_answer(eq, sym, None)
            return True
        else:
            return False

    def plug(self, eq):
        symt = eq.free_symbols.difference(self.answer.keys()).intersection(set(self.symbols))  # 和[x,y]取交集
        sym = symt.pop()
        known_sym, known_solve = zip(*self.answer.items())
        self.answer = defaultdict(list)
        exprs = []
        for known in zip(*known_solve):  # x, 2 or (x,z),[1,3]
            known = dict([(s, a) for s, a in zip(known_sym, known)])
            f = eq.subs(known)
            self.get_answer(f, sym, known)
            exprs.append(f)

        return exprs

    def plug_simple(self, eq):
        known_sym, known_solve = list(self.answer.items())[0]
        f = eq.subs({known_sym: known_solve[0]})
        return f

    def xiaoyuan(self, exprs):
        coeffs = [self.get_coeff(expr) for expr in exprs]
        syms = set.intersection(*[set(coeff.keys()) for coeff in coeffs])
        if len(syms) == 0:
            return None
        # max_com_temp = [(abs(lcm(*[coeff[sym] for coeff in coeffs])), sym) for sym in syms]
        # 把最小公倍数拆成带符号的和不带符号的 两部分
        lcm_lst_num = []  # 数值
        lcm_lst_sym = []  # 参数
        for sym in syms:
            lcm_of_sym = lcm(*[coeff[sym] for coeff in coeffs])
            if lcm_of_sym.is_Number:
                lcm_lst_num.append((abs(lcm_of_sym), sym))
            else:
                lcm_lst_sym.append((lcm_of_sym, sym))
        if len(lcm_lst_num) > 0:  # 存在系数是纯数字, 从中选出一个lcm最小的, 统一对应符号的系数
            max_com = sorted(lcm_lst_num, key=lambda x: x[0])
            flag_sym = False  #
        elif len(lcm_lst_sym) > 0:  # 如果系数都含符号, 就随便选一个
            max_com = lcm_lst_sym
            flag_sym = True  # flag表示两个系数都含符号

        # max_com = sorted([(abs(lcm(*[coeff[sym] for coeff in coeffs])), sym) for sym in syms], key = lambda x:x[0])
        lcm_num, use = max_com[0]
        if flag_sym:
            exprs = [(expr * (lcm_num / coeff[use])).expand() for expr, coeff in zip(exprs, coeffs)]
        else:
            exprs = [expr * abs(lcm_num / coeff[use]) for expr, coeff in zip(exprs, coeffs)]
        self.append(MULTIPLE % use, [(expr, zero) for expr in exprs])

        add_expr = exprs[0] + exprs[1]
        minus_expr = exprs[0] - exprs[1]

        if not add_expr.has(use):
            f = add_expr
            self.append(ADD % (self.index - 2, self.index - 1), [(f, zero)], False)
        elif not minus_expr.has(use):
            f = minus_expr
            self.append(MINUS % (self.index - 2, self.index - 1), [(f, zero)], False)
        else:
            return None

        if f.is_zero:
            raise ValueError("无限解。")
        elif f.is_number:
            raise ValueError("无解")
        return f

    def dairu(self):
        # select one equation
        free_symbols = array([len(eq.free_symbols) for eq in self.eqs])
        max_pows = {}
        for sym in self.symbols:
            max_pows[sym] = [get_pow(eq.lhs - eq.rhs, sym) for eq in self.eqs]

        pow_nums = [pow_num for _, pow_num in max_pows.items()]  # x:[1,2,3] y:[3,4,5] z:[6,7,8]
        pow_nums = array([max(t) for t in zip(*pow_nums)])

        coeffs = [self.get_coeff(eq.lhs - eq.rhs) for eq in self.eqs]  # [{},{},{}]
        coeffs_num = array([max([abs(n) for _, n in coeff.items()]) if coeff else 10 for coeff in coeffs])

        metrics = free_symbols * 100 + pow_nums * 10 + coeffs_num
        i = argsort(metrics)[0]

        # dairu
        eqs = []
        eq = self.eqs[i]

        if self.eq_num >= 3:
            candidates = [sym for sym, pows in max_pows.items() if max(pows) == 1]
            coeff = [(cand, abs(coeffs[i][cand])) for cand in candidates if cand in coeffs[i]]
            if len(coeff) == 0:
                raise ValueError("Too complex!")
            sym = sorted(coeff, key=lambda x: x[1])[0][0]
        else:
            sym = [k for k, v in coeffs[i].items() if k.is_Symbol][0]

        ans = self.get_answer(eq, sym, None)
        if self.eq_num >= 3 and len(ans.args) > 1:
            raise ValueError("Too complex!")  # multiple answer for the select equation

        last_eq = Eq(sym, self.answer[sym][0])
        self.append(GET % (i + self.index - 3), [last_eq])

        last_index = self.index - 1

        for t in argsort(metrics)[1:]:
            f = self.plug_simple(self.eqs[t])
            self.append(PLUG % (self.index - 1, t + self.index - 4), [f], False)
            eqs.append(f)

        return eqs, last_eq, last_index

    def _solve(self):
        is_simple = [self.check_simple(eq) for eq in self.eqs]  # is_simple是做什么的, 先不管
        if all(is_simple):
            # self.append("", self.eqs, False)
            return self
        elif any(is_simple):
            self.append("", self.eqs)
            eq_i = is_simple.index(True)
            n_i = is_simple.index(False)
            sym, ans = list(self.answer.items())[0]
            self.append(GET % (eq_i + self.index - 2), [(sym, a) for a in ans], False)
            last_eq = self.eqs[n_i]
            last_index = n_i + self.index - 2
        else:
            exprs = [Add(eq.lhs, -eq.rhs).expand() for eq in self.eqs]
            exprs = [self.factor(expr).together().cancel() for expr in exprs]
            temp = [get_all_child(expr, lambda x: x.is_Pow and x.args[1] == -1) for expr in exprs]
            if any(temp):
                for i, child in enumerate(temp):
                    if child:
                        exprs[i] = exprs[i] * (1 / Mul._from_args(child))

            if any([expr != eq.lhs for expr, eq in zip(exprs, self.eqs)]):
                self.append("", self.eqs, False)
                self.append(ZHENGLI, [(expr, zero) for expr in exprs])
            else:
                self.append("", self.eqs)

            last_eq = Eq(exprs[0], zero)
            last_index = self.index - 2

            f = self.xiaoyuan(exprs)
            if not f:
                # raise UnsolvableError("Too complex!")
                self.steps = []
                self.index = self.init_index
                self.append("", self.eqs)
                eqs, last_eq, last_index = self.dairu()
                f = eqs[0]

            # sym = f.free_symbols.pop()
            symt = f.free_symbols.intersection(set(self.symbols))  # 和[x,y]取交集
            sym = symt.pop()
            self.answer = defaultdict(list)
            ans = self.get_answer(f, sym, None)
            self.append(JIEDE, [(sym, a) for a in ans], False)

        exprs = self.plug(last_eq)
        self.append(PLUG_2 % last_index, [expr for expr in exprs], False)

        return self

    def solve(self):
        try:
            self.check()  # check检查变量数是否超过方程组数, 在做含参情况时, 我们暂时先不check
            # 还是要check的, check里设定了self.symbols
            self._solve()
        except Exception:
            pass

        return self

    @staticmethod
    def print_latex(expr, index=""):
        if isinstance(expr, tuple):
            latex_str = latex(Eq(expr[0], expr[1]))
            if "mathrm" in latex_str:
                latex_str = "%s = %s" % expr
        else:
            latex_str = latex(expr)
        if index:
            latex_str += "\\quad(%d)" % index
        return latex_str

    def append(self, text, exprs, index=True):
        if len(self.steps) > 0 and exprs == self.steps[-1][1]:
            return

        if index:
            self.steps.append([text, exprs, [self.index + i for i in range(len(exprs))]])
        else:
            self.steps.append([text, exprs, []])
        if index:
            self.index += len(exprs)

    def printing(self):
        print_steps = []
        for items in self.steps:
            if len(items) == 2:  # already printed
                print_steps.append(items)
                continue

            text, exprs, indice = items
            if len(indice) == len(exprs):
                eqs = [self.print_latex(expr, index) for expr, index in zip(exprs, indice)]
            else:
                eqs = [self.print_latex(expr) for expr in exprs]

            if len(set(eqs)) > 1:
                eqs = "\\\\".join(eqs)
                # eqs = "$$\\left\\{\\begin{array}{c}%s\\end{array}\\right.$$" % eqs
                eqs = "$$\\begin{equation}\\begin{cases}%s\\end{cases}\\end{equation}$$" % eqs
            else:
                eqs = "$$%s$$" % eqs[0]
            print_steps.append([text, eqs])

        if self.answer:
            known_sym, known_solve = zip(*self.answer.items())
            answers = []
            for known in zip(*known_solve):
                known = [self.print_latex(Eq(s, a)) for s, a in zip(known_sym, known)]
                answer = "\\\\".join(known)
                answers.append("$\\begin{equation}\\begin{cases}%s\\end{cases}\\end{equation}$" % answer)

            print_steps.append(["", " ".join(answers)])
        return print_steps


class Step3Equations(Step2Equations):
    eq_num = 3

    def __init__(self, eqs):
        self.eqs = eqs
        self.symbols = None
        self.steps = []
        self.index = 1
        self.answer = defaultdict(list)

    def solve(self):
        self.check()
        self.append("", self.eqs)

        eqs, last_eq, last_index = self.dairu()
        solver = Step2Equations(eqs, self.index).solve()
        steps = solver.printing()
        if not solver.answer:
            self.steps.extend(steps)
            return self

        steps[0][0] = "解方程："
        steps[-1][0] = "解得："
        self.steps.extend(steps)
        self.answer = solver.answer

        f = self.plug(last_eq)
        self.append(PLUG_2 % last_index, f, False)
        return self


if __name__ == '__main__':
    pass
